#include "refine/smooth_b.h"
#include "refine/window_pass.h"
#include <limits>
#include <iostream>
#include <set>
#include "imglib/imageclass.h"
#include <imglib/bitmap.h>
#include <imglib/bmpmorph.h>
#include <gslmath/rls.h>
#include <gslmath/matarithm.h>

using std::clog;
using std::endl;
using std::flush;

namespace refine {

namespace smooth_b {

const params default_params = {2, 2, 1, 100, /*100,*/ 0.0f, false};

typedef grey_mask::mask::elem_type grey_mask_value;
const grey_mask_value grey_mask_max_value = std::numeric_limits<grey_mask_value>::max();	//255
const grey_mask_value grey_mask_min_value = std::numeric_limits<grey_mask_value>::min();	//0
const grey_mask_value grey_mask_avg_value = (grey_mask_max_value + grey_mask_min_value) / 2;//127

typedef unsigned uint;

template <typename T>
inline T sqr(T t)
	{ return t * t; }

const static int fore_idx = -1; //this value in the map represents strict foreground pixel with which no variable is associated
const static int back_idx = -2;

/// maps pixel of unknown region to variable number
class varmap
{
public:
	varmap(img::bitmap const & strict_fore, img::bitmap const & strict_back) 
		: map(strict_fore.dim())
	{
		num = 0;
		for (img::image<int>::iterator v = map.begin(); v != map.end(); ++v) {
			if (strict_fore[v]) {
				assert (!strict_back[v]);
				*v = fore_idx;
			}
			else if (strict_back[v])
				*v = back_idx;
			else
				*v = num++;
		}
	}
	/// the total number of variables
	uint count() const
		{ return num; }
	int operator[] (size_point const & p) const
		{ return map[p]; }
	bool variable(size_point const & p) const
		{ return map[p] >= 0; }
private:
	img::image<int> map;
	uint num;
};

/// linear relation of central pixel's alpha is found in the form:
/// alpha = 
///   a0 + 
///   ra * sum_i R_i a_i +
///   ga * sum_i G_i a_i +
///   ba * sum_i B_i a_i +
///   a * sum_i a_i
struct relation {
	float a0, ra, ga, ba, a;
};
typedef std::vector<relation> relations;

class relation_constructor
{
public:
	relation_constructor(byte_image const & im_, 
		varmap const & vmap_, relations & rels_, std::vector<byte_color> const & colors)
		: im(im_), vmap(vmap_), rels(rels_)
	{ 
		rels.resize(vmap.count());
		remove_all(); 
		constructions_count = 0; 
		AAv = gsl_matrix_view_array(&AA[0][0], 4, 4);
		Gv = gsl_matrix_view_array(&Gm[0][0], 4, 4);
		hv = gsl_vector_view_array(h, 4);
		Iv = gsl_vector_view_array(I, 4);

		gsl_matrix * C = gsl_matrix_alloc(colors.size(), 4);
		gsl_vector * d = gsl_vector_alloc(colors.size());
		for (uint i = 0; i < colors.size(); ++i) {
			gsl_matrix_set(C, i, 0, colors[i].r);
			gsl_matrix_set(C, i, 1, colors[i].g);
			gsl_matrix_set(C, i, 2, colors[i].b);
			gsl_matrix_set(C, i, 3, grey_mask_max_value);
			gsl_vector_set(d, i, colors[i].a);
		}
		rls.init(C, d);
		gsl_vector_free(d);
		gsl_matrix_free(C);
	}
	size_extent dim() const
		{ return im.dim(); }
	uint count() const
		{ return constructions_count; }
	bool selected(size_point const & p) const
		{ return vmap.variable(p); }
	void add(size_point const & p) {
		byte_color const & I = im[p];
		add_sample(I.r, I.g, I.b);
	}
	void remove(size_point const & p) {
		byte_color const & I = im[p];
		remove_sample(I.r, I.g, I.b);
	}
	void remove_all() {
		n = 0;
		R = G = B = 0;
		RR = RG = RB = GG = GB = BB = 0;
	}
	void process(size_point const & p, size_rect const &) {
		++constructions_count;
		byte_color const & I = im[p];
		process_core(rels[vmap[p]], I.r, I.g, I.b);
	}
private:
	void add_sample(uint r, uint g, uint b) {
		++n;
		R += r; 
		G += g; 
		B += b; 
		RR += r * r;
		RG += r * g;
		RB += r * b;
		GG += g * g;
		GB += g * b;
		BB += b * b;
	}
	void remove_sample(uint r, uint g, uint b) {
		--n;
		R -= r; 
		G -= g; 
		B -= b; 
		RR -= r * r;
		RG -= r * g;
		RB -= r * b;
		GG -= g * g;
		GB -= g * b;
		BB -= b * b;
	}
	void process_core(relation & rel, uint rc, uint gc, uint bc) {
		AA[0][0] = RR;
		AA[0][1] = AA[1][0] = RG;
		AA[0][2] = AA[2][0] = RB;
		AA[0][3] = AA[3][0] = grey_mask_max_value * R;
		AA[1][1] = GG;
		AA[1][2] = AA[2][1] = GB;
		AA[1][3] = AA[3][1] = grey_mask_max_value * G;
		AA[2][2] = BB;
		AA[2][3] = AA[3][2] = grey_mask_max_value * B;
		AA[3][3] = n * grey_mask_max_value * grey_mask_max_value;

		rls.find_dep(&AAv.matrix, &Gv.matrix, &hv.vector);

		I[0] = rc;
		I[1] = gc;
		I[2] = bc;
		I[3] = grey_mask_max_value;

		rel.a0 = 0; //:= scalar_product(I, h)
		for (uint i = 0; i < 4; ++i)
			rel.a0 += I[i] * h[i];

		vec(&hv.vector) = cmatT(&Gv.matrix) * cvec(&Iv.vector);
		rel.ra = h[0];
		rel.ga = h[1];
		rel.ba = h[2];
		rel.a = h[3] * grey_mask_max_value;
	}
private:
	byte_image const & im;
	varmap const & vmap;
	relations & rels;
	uint constructions_count;
private:
	uint n;
	uint R, G, B;
	uint RR, RG, RB, GG, GB, BB;
private:
	rls_Ab_dep rls;
	double AA[4][4], Gm[4][4], h[4], I[4];
	gsl_matrix_view AAv, Gv;
	gsl_vector_view hv, Iv;
};

//template <typename Raster, typename Matte>
class smoother
{
public:
	smoother(byte_image const & src_, varmap const & vmap_, relations const & rels_,
		grey_mask::mask const & matte_, grey_mask::mask & next_matte_,
		img::bitmap & modifiable_, /*float max_condition_number_,*/ float sharpness_)
		: src(src_), vmap(vmap_), rels(rels_), matte(matte_), next_matte(next_matte_), modifiable(modifiable_),
		  /*max_condition_number(max_condition_number_),*/ sharpness(sharpness_)
		{ clear(); }
	void clear() {
		checks = achanges = 0;
		remove_all();
	}
	size_extent dim() const
		{ return src.dim(); }
	uint check_count() const
		{ return checks; }
	uint change_count() const
		{ return achanges; }
	bool selected(size_point const & p) const
		{ return modifiable[p]; }
	void add(size_point const & p) {
		byte_color const & I = src[p];
		add_sample(I.r, I.g, I.b, p.x, p.y, matte[p]);
	}
	void remove(size_point const & p) {
		byte_color const & I = src[p];
		remove_sample(I.r, I.g, I.b, p.x, p.y, matte[p]);
	}
	void remove_all() 
		{ n = RA = GA = BA = XA = YA = A = AA = 0; }
	///checks max_condition_number and returns true in the case of fail
	/*bool bad_condition_number() const {
		if (n == 0)
			return true;
		double a = AA;
		double b = grey_mask_max_value * A - AA;
		double c = grey_mask_max_value * (grey_mask_max_value * n - 2 * A) + AA;
		double D = sqr(a - c) + 4 * b * b;
		return sqr(max_condition_number + 1) * D > sqr((max_condition_number - 1) * (a + c));
	}*/
	void process(size_point const & p, size_rect const & wnd) {
		++checks;
		/*if (bad_condition_number()) {
			modifiable[p] = false;
			return;
		}*/

		relation const & r = rels[vmap[p]];
		float a = r.a0 + r.ra * RA + r.ga * GA + r.ba * BA + r.a * A + 0.5f; //0.5 for proper rounding

		if (sharpness > 0) {
			int sum_a = matte[p];
			int count = 1;
			if (p.x > 0) {
				sum_a += matte(p.x-1, p.y);
				++count;
			}
			if (p.x + 1 < dim().sx) {
				sum_a += matte(p.x+1, p.y);
				++count;
			}
			if (p.y > 0) {
				sum_a += matte(p.x, p.y-1);
				++count;
			}
			if (p.y + 1 < dim().sy) {
				sum_a += matte(p.x, p.y+1);
				++count;
			}
			a += (sharpness * 2 * sum_a) / count - sharpness * grey_mask_max_value;
		}

		grey_mask_value g;
		if (a <= grey_mask_min_value)
			g = grey_mask_min_value;
		else if (a >= grey_mask_max_value)
			g = grey_mask_max_value;
		else 
			g = grey_mask_value(a);

		bool not_greater_exists = false;
		bool not_less_exists = false;
		size_point i;
		for (i.y = wnd.y1; i.y < wnd.y2; ++i.y)
			for (i.x = wnd.x1; i.x < wnd.x2; ++i.x) {
				if (!not_greater_exists && matte[i] <= g)
					not_greater_exists = true;
				if (!not_less_exists && matte[i] >= g)
					not_less_exists = true;
				if (not_greater_exists && not_less_exists)
					break;
			}

		if (not_greater_exists && not_less_exists && matte[p] != g) {
			++achanges;
			next_matte[p] = g;
		}
		else
			modifiable[p] = false;
	}
private:
	void add_sample(uint r, uint g, uint b, uint x, uint y, uint a) {
		++n;
		RA += r * a;
		GA += g * a;
		BA += b * a;
		XA += x * a;
		YA += y * a;
		A += a;
		AA += a * a;
	}
	void remove_sample(uint r, uint g, uint b, uint x, uint y, uint a) {
		--n;
		RA -= r * a;
		GA -= g * a;
		BA -= b * a;
		XA -= x * a;
		YA -= y * a;
		A -= a;
		AA -= a * a;
	}
private:
	byte_image const & src;
	varmap const & vmap;
	relations const & rels;
	grey_mask::mask const & matte;
	grey_mask::mask & next_matte;
	img::bitmap & modifiable;
	//float max_condition_number;
	float sharpness;
	uint checks, achanges;
private:
	uint n, RA, GA, BA, XA, YA, A, AA;
};

void perform_core(
	const byte_image & src, 
	grey_mask::mask & matte, 
	const img::bitmap & strict_fore, 
	const img::bitmap & strict_back,
	img::bitmap & modifiable, ///< map of pixels that may be changed during the following pass
	const params & p)
{
	if (!p.silent)
		clog << "Volosok-B matting..." << endl;

	size_extent dim = src.dim();
	assert (dim == matte.dim());
	assert (dim == modifiable.dim());
	int_rect frame(0, 0, dim.sx, dim.sy);

	varmap vmap(strict_fore, strict_back);
	relations rels(vmap.count());
	window_pass(relation_constructor(src, vmap, rels, p.colors), p.whs);

	grey_mask::mask next_matte = matte;

	for (uint pass = 0; pass < p.passes_num; ++pass) {
		if (!p.silent)
			clog << "pass " << pass+1 << '/' << p.passes_num << flush;

		img::combine(modifiable, strict_fore, img::bit_and_not());
		img::combine(modifiable, strict_back, img::bit_and_not());

		smoother sm(src, vmap, rels, matte, next_matte, modifiable, /*p.max_condition_number,*/ p.sharpness);
		window_pass(sm, p.whs);
		if (!p.silent)
			clog << ": checked = " << sm.check_count() << ", alpha changes = " << sm.change_count() << endl;
		if (sm.change_count() == 0)
			break;

		matte = next_matte;

		img::dilate8(modifiable, p.whs);
	}
}


void perform(
	const byte_image & src, 
	const bit_mask::mask & m, 
	grey_mask::mask & matte, 
	const params & p)
{
	size_extent dim = src.dim();
	assert (dim == m.dim());

	matte.resize(dim);
	for (grey_mask::mask::iterator i = matte.begin(); i != matte.end(); ++i)
		*i = m[i] ? grey_mask_max_value : grey_mask_min_value;

	img::bitmap strict_fore, strict_back;
	img::erode4(strict_fore, m, p.fore_depth);
	img::dilate4(strict_back, m, p.back_depth);
	strict_back.inverse();

	img::bitmap modifiable = m; 
	{
		uint w = 1 + p.whs;// * (p.max_condition_number - 1) / (p.max_condition_number + 1);
		img::dilate4(modifiable, w);
		img::bitmap tmp = m;
		tmp.inverse();
		img::dilate4(tmp, w);
		img::combine(modifiable, tmp, img::bit_and());
	}

	perform_core(src, matte, strict_fore, strict_back, modifiable, p);
}

void perform(
	const byte_image & src, 
	grey_mask::mask & matte, 
	const params & p)
{
	size_extent dim = matte.dim();

	img::bitmap not_min(dim), not_max(dim), strict_fore(dim), strict_back(dim);
	for (grey_mask::mask::const_iterator i = matte.begin(); i != matte.end(); ++i) {
		not_min[i] = *i != grey_mask_min_value;
		not_max[i] = *i != grey_mask_max_value;
		strict_fore[i] = *i > grey_mask_avg_value;
	}

	strict_back = strict_fore;
	img::erode4(strict_fore, p.fore_depth);
	img::dilate4(strict_back, p.back_depth);
	strict_back.inverse();

	img::bitmap modifiable;
	{
		uint w = 1 + p.whs;// * (p.max_condition_number - 1) / (p.max_condition_number + 1);
		img::dilate4(modifiable, not_min, w);
		img::bitmap tmp;
		img::dilate4(tmp, not_max, w);
		img::combine(modifiable, tmp, img::bit_and());
	}

	perform_core(src, matte, strict_fore, strict_back, modifiable, p);
}

void perform(
	argb_byte_image & src,
	const params & p)
{
	//!!! temporary solution
	byte_image src2 = src; 
	grey_mask::mask matte = img::projection(src, &argb_byte_color::a);

	perform(src2, matte, p);

	img::projection(src, &argb_byte_color::a) = matte;
}

} //namespace smooth_b

} //namespace refine
